from flask import Flask, request, jsonify, render_template_string
from langfuse import Langfuse
import requests
import os
import uuid
import time
import re
from typing import Any, Dict, List, Union

app = Flask(__name__)

# 配置提示词名称与可选版本
PROMPT_NAME = os.getenv("LANGFUSE_PROMPT_NAME", "八号仓简单点的")
PROMPT_VERSION_DEFAULT = os.getenv("LANGFUSE_PROMPT_VERSION")  # 为空则取最新
# Dify 应用中用于承载系统提示词的变量名（需要与你在 Dify 模板里使用的占位符一致）
DIFY_SYSTEM_PROMPT_KEY = os.getenv("LANGFUSE_DIFY_SYSTEM_PROMPT_KEY", "system_prompt")

# 初始化 Langfuse
langfuse_client = Langfuse(
    public_key="pk-lf-ea07ebae-8acb-4e55-8c74-10615512b4c3",
    secret_key="sk-lf-4432bb59-ef86-47ed-a66e-ea916cbd8368",
    host="https://langfuse-dev.chinabeego.com",
    sample_rate=1.0,
    environment="default"
)

# 测试Langfuse连接
try:
    with langfuse_client.start_as_current_span(name="test_connection") as test_span:
        test_span.update(metadata={"test": "connection"})
    print("Langfuse连接测试成功")
except Exception as e:
    print(f"Langfuse连接测试失败: {e}")

# Dify 配置
DIFY_API_KEY = "app-ZPDhaSBS2xZcoOupEp3cPXnX"
DIFY_API_URL = "https://dify-pre.chinabeego.com/v1"

# 内存存储
conversation_sessions = {}

# 安全渲染 Langfuse 提示词（兼容旧 SDK）
def _render_prompt(prompt_obj: Any, variables: Dict[str, Any]) -> str:
    try:
        return prompt_obj.compile(**variables)
    except TypeError:
        pass
    except Exception:
        pass
    raw = prompt_obj.prompt
    if isinstance(raw, list):
        text_parts: List[str] = []
        for msg in raw:
            content = msg.get("content", "") if isinstance(msg, dict) else str(msg)
            text_parts.append(content)
        text = "\n".join(text_parts)
    else:
        text = str(raw)
    for k, v in (variables or {}).items():
        text = text.replace(f"{{{{{k}}}}}", str(v))
    return text

def generate_chat_response_with_prompt(user_message, session_id, mobile=None, prompt_version: str | None = None, system_prompt_key: str | None = None):
    """使用 Langfuse 提示词管理调用 Dify API 生成回复"""
    
    with langfuse_client.start_as_current_generation(
        name="dify_api_call_with_prompt",
        model="dify-ai-agent",
        input=user_message,
        metadata={
            "provider": "dify",
            "user_input": user_message,
            "session_id": session_id,
            "prompt_name": PROMPT_NAME,
            "prompt_version": prompt_version or PROMPT_VERSION_DEFAULT or "latest"
        }
    ) as generation:

        # 从 Langfuse 获取提示词（可指定版本）
        try:
            if prompt_version or PROMPT_VERSION_DEFAULT:
                prompt_obj = langfuse_client.get_prompt(name=PROMPT_NAME, version=prompt_version or PROMPT_VERSION_DEFAULT)
            else:
                prompt_obj = langfuse_client.get_prompt(name=PROMPT_NAME)
            print(f"成功获取提示词: {prompt_obj.name} v{getattr(prompt_obj, 'version', 'latest')}")
            # 将提示词对象关联到当前 generation（两种写法都支持不同 SDK 版本）
            try:
                langfuse_client.update_current_generation(prompt=prompt_obj)
            except Exception:
                pass
                generation.update(prompt=prompt_obj)
            
            if not mobile:
                match = re.search(r'1[3-9]\d{9}', user_message)
                mobile = match.group() if match else ""
            compiled_prompt = _render_prompt(prompt_obj, {"mobile": mobile, "user_input": user_message})
            print(f"提示词渲染成功，长度: {len(compiled_prompt)}")
            
        except Exception as e:
            print(f"获取或渲染提示词失败: {e}")
            compiled_prompt = "请以商场客服小8的身份，按规范为顾客服务。"

        headers = {
            "Authorization": f"Bearer {DIFY_API_KEY}",
            "Content-Type": "application/json"
        }
    
        inputs = {}
        inputs["mobile"] = mobile if mobile else "未提供手机号"
        # 选择本次请求使用的系统提示变量名
        key = (system_prompt_key or DIFY_SYSTEM_PROMPT_KEY)
        inputs[key] = compiled_prompt
        
        payload = {
            "inputs": inputs,
            "query": user_message,
            "response_mode": "streaming",
            "user": session_id
        }

        response_data = None
        try:
            print(f"发送请求到: {DIFY_API_URL}/chat-messages")
            print(f"请求头: {headers}")
            print(f"请求体: {payload}")
            
            response = requests.post(f"{DIFY_API_URL}/chat-messages", json=payload, headers=headers, timeout=120, stream=True)
            print(f"响应状态码: {response.status_code}")
            print(f"响应头: {dict(response.headers)}")
            
            response.raise_for_status()
            
            assistant_reply = ""
            response_data = {"answer": "", "conversation_id": "", "message_id": "", "trace_id": ""}
            
            for line in response.iter_lines():
                if line:
                    line_str = line.decode('utf-8')
                    if line_str.startswith('data: '):
                        data_str = line_str[6:]
                        if data_str.strip() == '[DONE]':
                            break
                        try:
                            data = eval(data_str)
                            if isinstance(data, dict):
                                if 'answer' in data:
                                    assistant_reply += data['answer']
                                    response_data['answer'] = assistant_reply
                                if 'conversation_id' in data:
                                    response_data['conversation_id'] = data['conversation_id']
                                if 'message_id' in data:
                                    response_data['message_id'] = data['message_id']
                                if 'id' in data:
                                    response_data['trace_id'] = data['id']
                        except:
                            continue
            
            if not assistant_reply:
                assistant_reply = "抱歉，无法获取AI回复。"

            estimated_prompt_tokens = len(user_message) // 4
            estimated_completion_tokens = len(assistant_reply) // 4
            usage_info = {
                "prompt_tokens": estimated_prompt_tokens,
                "completion_tokens": estimated_completion_tokens,
                "total_tokens": estimated_prompt_tokens + estimated_completion_tokens
            }

            generation.update(
                output=assistant_reply,
                usage=usage_info,
                metadata={
                    "dify_response": response_data,
                    "prompt_used": PROMPT_NAME,
                    "prompt_version": getattr(prompt_obj, 'version', prompt_version or PROMPT_VERSION_DEFAULT or 'latest'),
                    "dify_system_prompt_key": key,
                    "system_prompt_len": len(compiled_prompt)
                }
            )

            return {
                "answer": assistant_reply,
                "conversation_id": response_data.get('conversation_id', ''),
                "message_id": response_data.get('message_id', ''),
                "trace_id": response_data.get('trace_id', ''),
                "usage": usage_info
            }

        except requests.exceptions.RequestException as e:
            print(f"请求失败: {e}")
            generation.update(
                output="抱歉，服务暂时不可用，请稍后再试。",
                usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
                metadata={"error": str(e)}
            )
            return {
                "answer": "抱歉，服务暂时不可用，请稍后再试。",
                "conversation_id": "",
                "message_id": "",
                "trace_id": "",
                "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
            }

@app.route('/')
def index():
    """主页"""
    return render_template_string('''
    <!DOCTYPE html>
    <html>
    <head>
        <title>Langfuse + Dify 提示词管理演示</title>
        <meta charset="utf-8">
        <style>
            body { font-family: Arial, sans-serif; margin: 40px; }
            .container { max-width: 800px; margin: 0 auto; }
            .form-group { margin-bottom: 20px; }
            label { display: block; margin-bottom: 5px; font-weight: bold; }
            input, textarea, select { width: 100%; padding: 10px; border: 1px solid #ddd; border-radius: 4px; }
            button { background-color: #007bff; color: white; padding: 10px 20px; border: none; border-radius: 4px; cursor: pointer; }
            button:hover { background-color: #0056b3; }
            .response { margin-top: 20px; padding: 15px; background-color: #f8f9fa; border-radius: 4px; }
            .error { color: red; }
            .success { color: green; }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🤖 Langfuse + Dify 提示词管理演示</h1>
            <p>使用 Langfuse 管理提示词，通过 Dify API 进行 AI 对话</p>
            
            <form id="chatForm">
                <div class="form-group">
                    <label for="user_message">用户消息:</label>
                    <textarea id="user_message" name="user_message" rows="3" placeholder="请输入您的问题..." required></textarea>
                </div>
                
                <div class="form-group">
                    <label for="mobile">手机号 (可选):</label>
                    <input type="text" id="mobile" name="mobile" placeholder="请输入手机号">
                </div>
                
                <div class="form-group">
                    <label for="prompt_version">提示词版本 (可选):</label>
                    <input type="text" id="prompt_version" name="prompt_version" placeholder="例如: 2 或 1.0.1">
                </div>
                
                <div class="form-group">
                    <label for="system_prompt_key">Dify系统提示变量名 (可选):</label>
                    <input type="text" id="system_prompt_key" name="system_prompt_key" placeholder="默认 system_prompt">
                </div>
                
                <button type="submit">发送消息</button>
            </form>
            
            <div id="response" class="response" style="display: none;">
                <h3>AI 回复:</h3>
                <div id="response_content"></div>
            </div>
        </div>

        <script>
            document.getElementById('chatForm').addEventListener('submit', async function(e) {
                e.preventDefault();
                
                const formData = new FormData(this);
                const responseDiv = document.getElementById('response');
                const responseContent = document.getElementById('response_content');
                
                responseDiv.style.display = 'block';
                responseContent.innerHTML = '正在处理中...';
                
                try {
                    const response = await fetch('/chat', {
                        method: 'POST',
                        headers: {
                            'Content-Type': 'application/json',
                        },
                        body: JSON.stringify({
                            user_message: formData.get('user_message'),
                            mobile: formData.get('mobile'),
                            prompt_version: formData.get('prompt_version'),
                            system_prompt_key: formData.get('system_prompt_key')
                        })
                    });
                    
                    const data = await response.json();
                    
                    if (data.success) {
                        responseContent.innerHTML = `
                            <p><strong>回复:</strong> ${data.answer}</p>
                            <p><strong>会话ID:</strong> ${data.conversation_id}</p>
                            <p><strong>消息ID:</strong> ${data.message_id}</p>
                            <p><strong>Token使用:</strong> ${data.usage.total_tokens} (输入: ${data.usage.prompt_tokens}, 输出: ${data.usage.completion_tokens})</p>
                        `;
                    } else {
                        responseContent.innerHTML = `<p class="error">错误: ${data.error}</p>`;
                    }
                } catch (error) {
                    responseContent.innerHTML = `<p class="error">请求失败: ${error.message}</p>`;
                }
            });
        </script>
    </body>
    </html>
    ''')

@app.route('/chat', methods=['POST'])
def chat():
    """聊天接口"""
    try:
        data = request.get_json()
        user_message = data.get('user_message', '').strip()
        mobile = data.get('mobile', '')
        prompt_version = data.get('prompt_version') or PROMPT_VERSION_DEFAULT
        system_prompt_key = (data.get('system_prompt_key') or '').strip() or DIFY_SYSTEM_PROMPT_KEY
        
        if not user_message:
            return jsonify({
                "success": False,
                "error": "用户消息不能为空"
            })
        
        session_id = str(uuid.uuid4())
        
        result = generate_chat_response_with_prompt(
            user_message=user_message,
            session_id=session_id,
            mobile=mobile,
            prompt_version=prompt_version,
            system_prompt_key=system_prompt_key
        )
        
        return jsonify({
            "success": True,
            "answer": result["answer"],
            "conversation_id": result["conversation_id"],
            "message_id": result["message_id"],
            "trace_id": result["trace_id"],
            "usage": result["usage"]
        })
        
    except Exception as e:
        print(f"聊天接口错误: {e}")
        return jsonify({
            "success": False,
            "error": str(e)
        })

@app.route('/health', methods=['GET'])
def health():
    """健康检查接口"""
    return jsonify({
        "status": "healthy",
        "langfuse": "connected",
        "dify": "configured",
        "timestamp": time.time()
    })

@app.route('/prompts', methods=['GET'])
def list_prompts():
    """列出所有提示词"""
    try:
        prompts = langfuse_client.get_prompts()
        prompt_list = []
        for prompt in prompts:
            prompt_list.append({
                "name": prompt.name,
                "version": prompt.version,
                "description": getattr(prompt, 'description', ''),
                "created_at": getattr(prompt, 'created_at', ''),
                "updated_at": getattr(prompt, 'updated_at', '')
            })
        
        return jsonify({
            "success": True,
            "prompts": prompt_list
        })
    except Exception as e:
        return jsonify({
            "success": False,
            "error": str(e)
        })

@app.route('/prompts/<prompt_name>', methods=['GET'])
def get_prompt(prompt_name):
    """获取特定提示词"""
    try:
        version = request.args.get('version')
        if version:
            prompt = langfuse_client.get_prompt(name=prompt_name, version=version)
        else:
            prompt = langfuse_client.get_prompt(name=prompt_name)
        
        return jsonify({
            "success": True,
            "prompt": {
                "name": prompt.name,
                "version": prompt.version,
                "prompt": prompt.prompt,
                "description": getattr(prompt, 'description', ''),
                "created_at": getattr(prompt, 'created_at', ''),
                "updated_at": getattr(prompt, 'updated_at', '')
            }
        })
    except Exception as e:
        return jsonify({
            "success": False,
            "error": str(e)
        })

if __name__ == '__main__':
    print("🚀 启动 Langfuse + Dify 提示词管理 API 服务")
    print("=" * 50)
    print("服务地址: http://localhost:5000")
    print("API 接口:")
    print("  POST /chat - 聊天接口")
    print("  GET  /prompts - 列出所有提示词")
    print("  GET  /prompts/<name> - 获取特定提示词")
    print("  GET  /health - 健康检查")
    print("=" * 50)
    
    app.run(debug=True, host='0.0.0.0', port=5000)
